import { parentPort } from 'worker_threads';
import crypto from 'crypto';

const port = parentPort;
if (!port) throw new Error('IllegalState');

const encoding: 'base64' | 'hex' = 'base64';

const keyLength = 'modp14';//密钥对长度

port.on('message', (msg) => {
  console.log('DHWorker 已加载');
  if (msg.method === 'getKeyPair') {
    const keyPair = getDHKeyPair();
    // 返回密钥对给主线程
    port.postMessage(keyPair);
  }
  if (msg.method === 'key-negotiation') {
    try {
      const shareKey = keyNegotiation(
        msg.dhPrivateKey,
        msg.dhPublicKey,
        msg.otherPublicKey
      );
      port.postMessage({ shareKey });
    } catch (error) {
      console.log('密钥协商错误：', error);
    }
  }
});

/**
 * 使用给定的私钥和公钥进行DH密钥协商，
 * 返回一个共享的秘密密钥。
 * @param {string} dhPrivateKey DH私钥。
 * @param {string} dhPublicKey DH公钥。
 * @param {string} otherPublicKey 对方的DH公钥。
 * @returns {string} 共享的秘密密钥。
 */
const keyNegotiation = (
  dhPrivateKey: string,
  dhPublicKey: string,
  otherPublicKey: string
): string => {
  
  //根据密钥信息生成 DiffieHellman 类的实例
  const temp = crypto.getDiffieHellman(keyLength).getPrime();
  const dh = crypto.createDiffieHellman(temp);

  // 将密钥转换为 Buffer
  const dhPrivateKeyBuffer = Buffer.from(dhPrivateKey, encoding);
  const dhPublicKeyBuffer = Buffer.from(dhPublicKey, encoding);
  const otherPublicKeyBuffer = Buffer.from(otherPublicKey, encoding);

  // 重写 DiffieHellman 类的实例的私钥和公钥
  dh.setPrivateKey(dhPrivateKeyBuffer);
  dh.setPublicKey(dhPublicKeyBuffer);

  // 计算共享密钥
  const shareKey = dh.computeSecret(otherPublicKeyBuffer);

  return crypto.createHash('sha256').update(shareKey).digest('hex');
};


/**
 * 生成给定长度的Diffie-Hellman密钥对。
 * @param {512 | 1024 | 2048 | 4096} length 密钥对的长度。
 * @returns {Object} Diffie-Hellman实例和PEM格式的私钥和公钥。
 */
const getDHKeyPair = () => {
  const dh = crypto.getDiffieHellman(keyLength);
  dh.generateKeys();
  return {
    dh: dh,
    publicKey: dh.getPublicKey(encoding),
    privateKey: dh.getPrivateKey(encoding)
  };
};
